import pybullet as p
import pybullet_data
import time
import math
import matplotlib.pyplot as plt

from numpy.lib.npyio import savez_compressed


import numpy as np

class ecoRobot():
    def __init__(self):
        self.physics = p.connect(p.GUI)
        p.setGravity(0,0,-9.81,self.physics)
        p.setAdditionalSearchPath(pybullet_data.getDataPath())
        p.loadURDF("plane.urdf")
        self.robot_id = None
        self.wheel_indices = [0,1]
        self.time_step = 1./240.
        self.x_increment = 0.01
        self.amplitude = 1.0
        self.frequency = 1.0
        self.scale = 1.0
        self.len_of_robot = 0.2
        self.wid_of_robot = 0.16
        self.radius = 0.025
        self.fig = plt.figure()
        robot_start_position = [0,0,0.026]
        robot_start_orientation = p.getQuaternionFromEuler([0,0,0])
        self.robot_id = p.loadURDF("/home/maxwene/project/code/data_structure_learning/强化学习/simulator/robot/x2.urdf",robot_start_position,robot_start_orientation)
        for j in range(p.getNumJoints(self.robot_id )):
            print(j,p.getJointInfo(self.robot_id, j))
    def mov(self,v,w):
        wheel_right_front = -(v+(w*(self.wid_of_robot )/2))/self.radius
        wheel_left_front  = -(v-(w*(self.wid_of_robot )/2))/self.radius

        return [wheel_left_front,wheel_right_front]

    def mov_robot(self,x):
        
        mov_list = self.mov(x[3],x[4])
        v = x[3]
        w = x[4]
        traj_motion = []
        traj_pybullet = []
        location, quater = p.getBasePositionAndOrientation(self.robot_id)
        roll,pitch,yaw = p.getEulerFromQuaternion(quater)
        x = [location[0],location[1],yaw,0,0]
        print(x)
        print(location[0],location[1],yaw)
        for _ in range(10000):
            
            for wheel_index in self.wheel_indices:
                p.setJointMotorControl2(
                    self.robot_id,
                    wheel_index,
                    p.VELOCITY_CONTROL,
                    targetVelocity= mov_list[wheel_index],
                    force = 100
                )
            p.stepSimulation(self.physics)

            location, quater = p.getBasePositionAndOrientation(self.robot_id)
            roll,pitch,yaw = p.getEulerFromQuaternion(quater)
            
            x = self.motion(x,[v,w],1./120.)
            traj_motion.append(x.copy())
            print("motion",x)
            traj_pybullet.append([location[1],location[0],-yaw].copy())
            x = [location[1],location[0],-yaw,x[3],x[4]]
            
            
            print("imu",location[0],location[1],yaw)
            if abs(yaw) >= math.pi/2:
                break
            time.sleep(1./120.)
        
        self.show(traj_motion,"motion")
        self.show(traj_pybullet,"pybullet")
        plt.show()

    def reset(self):
        robot_start_orientation = p.getQuaternionFromEuler([0,0,0])
        p.resetBasePositionAndOrientation(self.robot_id,
                                            [0,0,robot_start_orientation[2]],
                                            robot_start_orientation)

    def close(self):
        p.disconnect()

    def motion(self,x,u,dt):
        x[2]+= u[1]*dt
        x[0]+= u[0]*math.cos(x[2])*dt
        x[1]+= u[0]*math.sin(x[2])*dt
        if(x[2] > math.pi*2):
            x[2] = x[2] - math.pi*2
        x[3] = u[0]
        x[4] = u[1]
        return x
    def show(self,traj,label):
        x = []
        y = []
        for i in traj:
            x.append(i[0])
            y.append(i[1])

        
        ax = self.fig.gca()
        figure = ax.plot(x,y,label=label)
        


def main():
    eRobot = ecoRobot()
    x = [0,0,0,1,1]
    eRobot.mov_robot(x)
    eRobot.close()

main()